package top.go2do.common.filter;

import org.slf4j.MDC;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

import javax.servlet.*;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.UUID;

/**
 * @author xingyuji
 * @date 2025/3/10 06:54
 * @description
 */
@Component
public class RequestIdFilter implements Filter {

    private static final String REQUEST_ID_HEADER = "X-Request-Id";
    public static final String REQUEST_ID_MDC = "requestId";
    private static final String DEVICE_ID_HEADER = "X-Device-Id";
    public static final String DEVICE_ID_MDC = "deviceId";

    @Override
    public void doFilter(ServletRequest request, ServletResponse response,
                         FilterChain chain) throws IOException, ServletException {
        HttpServletRequest httpServletRequest = (HttpServletRequest) request;
        // 尝试从请求头获取现有ID，否则生成新ID
        String requestId = httpServletRequest.getHeader(REQUEST_ID_HEADER);
        if (requestId == null || requestId.isEmpty()) {
            requestId = UUID.randomUUID().toString().replace("-", "");
        }

        // 存入MDC
        MDC.put(REQUEST_ID_MDC, requestId);

        // 尝试从请求头获取现有ID，否则生成新ID
        String deviceId = httpServletRequest.getHeader(DEVICE_ID_HEADER);
        if (StringUtils.hasLength(deviceId)) {
            // 存入MDC
            MDC.put(DEVICE_ID_MDC, deviceId);
        }

        try {
            HttpServletResponse httpServletResponse = (HttpServletResponse) response;
            // 将requestId添加到响应头（可选）
            httpServletResponse.addHeader(REQUEST_ID_HEADER, requestId);
            if (StringUtils.hasLength(deviceId)) {
                httpServletResponse.addHeader(DEVICE_ID_HEADER, deviceId);
            }
            chain.doFilter(request, response);
        } finally {
            // 确保请求结束后清理MDC
            MDC.clear();
        }
    }
}
